# 编写函数 扭曲和变换视角
# 扭曲和转换.py
import pickle
import cv2
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

#读取保存的相机矩阵 和 失真系数
# 方便后续使用 cv2.calibrateCamera() 方法
dist_pickle = pickle.load( open( "wide_dist_pickle.p", "rb" ) )
mtx = dist_pickle["mtx"]
dist = dist_pickle["dist"]

# 读入图片
img = cv2.imread('test_image/test_image.png')

cv2.imshow('imshow', img)
cv2.waitKey(0)


nx = 8 
ny = 6 

# def corners_unwarp(img, nx, ny, mtx, dist):
# 	# 1) 使用 mtx 和 dist 进行不失真
# 	undist = cv2.undistort(img, mtx, dist, None, mtx)
	
# 	# 2) 转换为灰度图
# 	gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  
# 	img_size = (gray.shape[1], gray.shape[0])
	
# 	# 3) 找到棋盘的角
# 	ret, corners = cv2.findChessboardCorners(gray, (nx, ny), None)
    
# 	# 4) 如果找到角: 
# 	if ret == True:
# 		# a) 画角
# 		cv2.drawChessboardCorners(img, (nx, ny), corners, ret)

# 		src = np.float32(
# 			[[580, 275],
# 			 [790, 230],
# 			 [580, 515],
# 			 [790, 505]])
    
# 		dst = np.float32(
# 			[[580, 275],
# 			 [790, 275],
# 			 [580, 515],
# 			 [790, 515]])

# 		M = cv2.getPerspectiveTransform(src, dst)


# 		warped = cv2.warpPerspective(undist, M, img_size)

# 	return warped, M

def corners_unwarp(img, nx, ny, mtx, dist):
    undist = cv2.undistort(img, mtx, dist, None, mtx)
    gray = cv2.cvtColor(undist, cv2.COLOR_BGR2GRAY)
    ret, corners = cv2.findChessboardCorners(gray, (nx, ny), None)

    if ret == True:
        cv2.drawChessboardCorners(undist, (nx, ny), corners, ret)
  
        offset = 100 
        
        img_size = (gray.shape[1], gray.shape[0])

        src = np.float32([corners[0], 
        				  corners[nx-1], 
        				  corners[-1], 
        				  corners[-nx]])

        dst = np.float32([[offset, offset], 
        				  [img_size[0]-offset, offset], 
                          [img_size[0]-offset, img_size[1]-offset], 
                          [offset, img_size[1]-offset]])

        M = cv2.getPerspectiveTransform(src, dst)
       
        warped = cv2.warpPerspective(undist, M, img_size)

    
    return warped, M

top_down, perspective_M = corners_unwarp(img, nx, ny, mtx, dist)
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(24, 9))
f.tight_layout()
ax1.imshow(img)
ax1.set_title('Original Image', fontsize=50)
ax2.imshow(top_down)
ax2.set_title('Undistorted and Warped Image', fontsize=50)
plt.subplots_adjust(left=0., right=1, top=0.9, bottom=0.)

plt.show()

r, g, b = cv2.split(top_down)
result = cv2.merge((b, g, r))

cv2.imwrite("test_image/output_test_image.jpg", result)